// Copyright 2014 The mqrouter Author. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package xkafka

import (
	"context"
	"encoding/json"
	"gitee.com/go-mid/infra/xmq/mqconf"
	"time"

	"gitee.com/go-mid/infra/xtime"
	"gitee.com/go-mid/infra/xtrace"
	"github.com/segmentio/kafka-go"
)

type WriterConfig struct {
	mqconf.MQConfig
	kafka.WriterConfig
}
type KafkaWriter struct {
	*kafka.Writer
	// NOTE: KafkaWriter 没有 config 的 getter，故在此保留一份
	config kafka.WriterConfig
	stop   chan struct{}
}

func NewKafkaWriter(writerConfig *WriterConfig) *KafkaWriter {
	writer := kafka.NewWriter(writerConfig.WriterConfig)
	kafkaWriter := &KafkaWriter{
		Writer: writer,
		config: writerConfig.WriterConfig,
		stop:   make(chan struct{}),
	}
	go kafkaWriter.init()

	return kafkaWriter
}

func (m *KafkaWriter) init() {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			stats := m.Stats()
			statKafkaWriter(m.config.Topic, &stats)
		case <-m.stop:
			return
		}
	}
}

func (m *KafkaWriter) setKafkaSpanTags(span xtrace.Span) {
	config := m.config
	span.SetTag(xtrace.TagSpanKind, xtrace.SpanKindProducer)
	span.SetTag(xtrace.TagMessageBusDestination, config.Topic)
	span.SetTag(xtrace.TagMessagingSystem, xtrace.MessagingSystemKafka)
	span.SetTag(xtrace.TagMessagingDestinationKind, xtrace.MessagingDestinationKindTopic)
	span.SetTag(xtrace.TagMessagingDestination, config.Topic)

}
func (m *KafkaWriter) WriteMsg(ctx context.Context, k string, v interface{}) (string, error) {
	span := xtrace.SpanFromContext(ctx)
	if span != nil {
		m.setKafkaSpanTags(span)
	}

	msg, err := json.Marshal(v)
	if err != nil {
		return "", err
	}

	st := xtime.NewTimeStat()
	err = m.WriteMessages(ctx, kafka.Message{
		Key:   []byte(k),
		Value: msg,
	})
	mqconf.StatReqDuration(ctx, m.config.Topic, "KafkaWriter.WriteMessages", mqconf.TraceMessageBusTypeKafka, st.Millisecond())

	return "", err
}

func (m *KafkaWriter) WriteMsgs(ctx context.Context, msgs []kafka.Message) error {
	span := xtrace.SpanFromContext(ctx)
	if span != nil {
		m.setKafkaSpanTags(span)
	}

	st := xtime.NewTimeStat()
	err := m.WriteMessages(ctx, msgs...)
	mqconf.StatReqDuration(ctx, m.config.Topic, "KafkaWriter.WriteMessages", mqconf.TraceMessageBusTypeKafka, st.Millisecond())

	return err
}

func (m *KafkaWriter) Close(ctx context.Context) (err error) {
	m.stop <- struct{}{}
	return m.Writer.Close()
}
